package com.wfh.mianshiji.ai.chatmemory;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.wfh.mianshiji.model.entity.Conversion;
import com.wfh.mianshiji.model.entity.Interview;
import com.wfh.mianshiji.service.ConversionService;
import com.wfh.mianshiji.service.InterviewService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

/**
 * @Title: DbChatMemory
 * @Author wangfenghuan
 * @Package com.wfh.mianshiji.ai.chatmemory
 * @Date 2025/10/31 18:49
 * @description:
 */
@Component
@Slf4j
public class InterviewDbChatMemory implements ChatMemory {

    @Resource
    @Lazy
    private ConversionService conversionService;

    @Resource
    @Lazy
    private InterviewService interviewService;

    @Override
    public void add(@NotNull String conversationId, @NotNull List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        saveConversation(conversationId, messages);
    }

    /**
     * 获取指定会话的最近 N 条消息（自定义扩展方法）
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        int fromIndex = Math.max(0, messageList.size() - lastN);
        return new ArrayList<>(messageList.subList(fromIndex, messageList.size()));
    }

    /**
     * 保存消息到当前会话id
     *
     * @param conversationId
     * @param messages
     */
    public void saveConversation(String conversationId, List<Message> messages) {
        Interview interview = interviewService.getOne(new QueryWrapper<>(Interview.class).eq("id", conversationId));
        Long userId = interview.getUserId();
        for (Message message : messages) {
            // 获取当前消息的类型
            String type = message.getMessageType().getValue();
            Conversion conversion = new Conversion();
            conversion.setInterviewId(Long.valueOf(conversationId));
            String text = message.getText();
            conversion.setUserId(userId);
            if (type.equals(MessageType.ASSISTANT.getValue())){
                conversion.setMessageType("ai");
                conversion.setMessage(text);
            }else if (type.equals(MessageType.USER.getValue())){
                conversion.setMessage(text);
                conversion.setMessageType("user");
            }else {
                log.warn("未知消息类型");
            }
           // 插入一条消息记录
            conversionService.save(conversion);
        }
    }

    /**
     * 获取或创建会话
     *
     * @param conversationId
     * @return
     */
    public List<Message> getOrCreateConversation(String conversationId) {
        List<Conversion> interviews = conversionService.getBaseMapper().selectList(new QueryWrapper<>(Conversion.class).eq("interviewId", conversationId));
        // 会话不为空，代表会话存在
        List<Message> messageList = new ArrayList<>();
        if (!interviews.isEmpty()) {
            for (Conversion conversion : interviews) {
                String messageType = conversion.getMessageType();
                String message = conversion.getMessage();
                if (message == null){
                    continue;
                }
                // 如果是用户消息
                if ("user".equals(messageType)){
                    messageList.add(new UserMessage(message));
                } else if ("ai".equals(messageType)) {
                    messageList.add(new AssistantMessage(message));
                }else {
                    log.warn("未知的消息类型");
                }
            }

        }
        return messageList;
    }

    /**
     * 清除指定会话的所有消息
     * @param conversationId
     */
    @Override
    public void clear(@NotNull String conversationId) {
        try {
            Long interviewId = Long.valueOf(conversationId);
            conversionService.getBaseMapper()
                    .delete(new QueryWrapper<Conversion>().eq("interviewId", interviewId));
        } catch (NumberFormatException e) {
            log.warn("无法清除会话，ID 非法: {}", conversationId);
        }
    }
}
